/*
 * 树算法模板库 - 树上的高效算法实现
 * 使用场景: 树形结构相关问题，如路径查询、子树操作、祖先查询等
 * 包含: 树上DFS、LCA、树的直径、重心、树上路径、换根DP等
 */

#pragma once
#include "basic_template.hpp"

// 树的基本表示 - 支持带权和无权树
// 使用场景: 作为其他树算法的基础数据结构
// 特点: 提供树的基本操作和属性查询
class Tree
{
public:
    int n, root;
    vector<vector<int>> adj;
    vector<vector<pair<int, int>>> weighted_adj;
    vector<int> parent, depth;
    vector<ll> dist_from_root;

    Tree(int vertices, int r = 0) : n(vertices), root(r)
    {
        adj.resize(n);
        weighted_adj.resize(n);
        parent.resize(n);
        depth.resize(n);
        dist_from_root.resize(n);
    }

    void add_edge(int u, int v, int w = 1)
    {
        adj[u].push_back(v);
        adj[v].push_back(u);
        weighted_adj[u].push_back({v, w});
        weighted_adj[v].push_back({u, w});
    }

    // 预处理：计算深度、父节点、到根的距离
    void preprocess()
    {
        fill(parent.begin(), parent.end(), -1);
        fill(depth.begin(), depth.end(), 0);
        fill(dist_from_root.begin(), dist_from_root.end(), 0);

        function<void(int, int)> dfs = [&](int u, int p) {
            parent[u] = p;
            for (auto [v, w] : weighted_adj[u]) {
                if (v != p) {
                    depth[v] = depth[u] + 1;
                    dist_from_root[v] = dist_from_root[u] + w;
                    dfs(v, u);
                }
            }
        };

        dfs(root, -1);
    }
};

// 自顶向下DFS - 从根到叶的信息传递
// 使用场景: 路径统计、节点标记、树的遍历
// 时间复杂度: O(n)
// 典型应用: 路径和、深度计算、子树大小
template <typename T>
class TopDownDFS
{
private:
    const Tree &tree;
    vector<T> result;

public:
    TopDownDFS(const Tree &t) : tree(t), result(t.n) {}

    // 自定义DFS处理函数
    template <typename ProcessFunc>
    void dfs(int u, int parent, T current_value, ProcessFunc process)
    {
        // 处理当前节点
        result[u] = process(u, current_value);

        // 递归处理子节点
        for (int v : tree.adj[u]) {
            if (v != parent) {
                dfs(v, u, result[u], process);
            }
        }
    }

    const vector<T> &get_result() const { return result; }
};

// 自底向上DFS - 从叶到根的信息汇聚
// 使用场景: 子树统计、最优化问题、树形DP
// 时间复杂度: O(n)
// 典型应用: 子树大小、子树和、最大路径
template <typename T>
class BottomUpDFS
{
private:
    const Tree &tree;
    vector<T> result;

public:
    BottomUpDFS(const Tree &t) : tree(t), result(t.n) {}

    // 自定义DFS处理函数
    template <typename ProcessFunc>
    T dfs(int u, int parent, ProcessFunc process)
    {
        vector<T> child_results;

        // 先递归处理所有子节点
        for (int v : tree.adj[u]) {
            if (v != parent) {
                child_results.push_back(dfs(v, u, process));
            }
        }

        // 基于子节点结果处理当前节点
        result[u] = process(u, child_results);
        return result[u];
    }

    const vector<T> &get_result() const { return result; }
};

// LCA (最近公共祖先) - 倍增算法实现
// 使用场景: 树上路径查询、距离计算、路径LCA
// 时间复杂度: 预处理O(n log n)，查询O(log n)
// 典型应用: 树上路径问题、树的查询、图论算法
class LCA
{
private:
    int n, LOG;
    vector<vector<int>> up;
    vector<int> depth;
    const Tree &tree;

public:
    LCA(const Tree &t) : tree(t), n(t.n)
    {
        LOG = 0;
        while ((1 << LOG) <= n)
            LOG++;
        up.assign(n, vector<int>(LOG, -1));
        depth.assign(n, 0);
        preprocess();
    }

    void preprocess()
    {
        function<void(int, int)> dfs = [&](int u, int p) {
            up[u][0] = p;
            for (int i = 1; i < LOG; i++) {
                if (up[u][i - 1] != -1) {
                    up[u][i] = up[up[u][i - 1]][i - 1];
                }
            }
            for (int v : tree.adj[u]) {
                if (v != p) {
                    depth[v] = depth[u] + 1;
                    dfs(v, u);
                }
            }
        };

        dfs(tree.root, -1);
    }

    int lca(int u, int v)
    {
        if (depth[u] < depth[v])
            swap(u, v);

        // 将u提升到与v相同的深度
        int diff = depth[u] - depth[v];
        for (int i = 0; i < LOG; i++) {
            if ((diff >> i) & 1) {
                u = up[u][i];
            }
        }

        if (u == v)
            return u;

        // 二分查找LCA
        for (int i = LOG - 1; i >= 0; i--) {
            if (up[u][i] != up[v][i]) {
                u = up[u][i];
                v = up[v][i];
            }
        }

        return up[u][0];
    }

    int distance(int u, int v)
    {
        return depth[u] + depth[v] - 2 * depth[lca(u, v)];
    }

    // 获取u到v路径上距离u为k的节点
    int kth_ancestor(int u, int k)
    {
        if (depth[u] < k)
            return -1;

        for (int i = 0; i < LOG; i++) {
            if ((k >> i) & 1) {
                u = up[u][i];
                if (u == -1)
                    return -1;
            }
        }
        return u;
    }
};

// 树的直径 - 找到树中最长的路径
// 使用场景: 树的性质分析、最优化问题
// 时间复杂度: O(n)
// 典型应用: 网络设计、树的中心、最长路径
class TreeDiameter
{
private:
    const Tree &tree;
    int diameter_length;
    pair<int, int> diameter_endpoints;
    vector<int> diameter_path;

    pair<int, int> bfs_farthest(int start)
    {
        vector<int> dist(tree.n, -1);
        queue<int> q;
        q.push(start);
        dist[start] = 0;

        int farthest = start;
        int max_dist = 0;

        while (!q.empty()) {
            int u = q.front();
            q.pop();

            for (int v : tree.adj[u]) {
                if (dist[v] == -1) {
                    dist[v] = dist[u] + 1;
                    q.push(v);

                    if (dist[v] > max_dist) {
                        max_dist = dist[v];
                        farthest = v;
                    }
                }
            }
        }

        return {farthest, max_dist};
    }

public:
    TreeDiameter(const Tree &t) : tree(t)
    {
        calculate_diameter();
    }

    void calculate_diameter()
    {
        // 第一次BFS：从任意点找到最远点
        auto [node1, _] = bfs_farthest(0);

        // 第二次BFS：从最远点找到直径的另一端
        auto [node2, length] = bfs_farthest(node1);

        diameter_length = length;
        diameter_endpoints = {node1, node2};

        // 构建直径路径
        vector<int> parent(tree.n, -1);
        vector<bool> visited(tree.n, false);
        queue<int> q;
        q.push(node1);
        visited[node1] = true;

        while (!q.empty()) {
            int u = q.front();
            q.pop();

            if (u == node2)
                break;

            for (int v : tree.adj[u]) {
                if (!visited[v]) {
                    visited[v] = true;
                    parent[v] = u;
                    q.push(v);
                }
            }
        }

        // 重建路径
        diameter_path.clear();
        int curr = node2;
        while (curr != -1) {
            diameter_path.push_back(curr);
            curr = parent[curr];
        }
        reverse(diameter_path.begin(), diameter_path.end());
    }

    int get_diameter() const { return diameter_length; }
    pair<int, int> get_endpoints() const { return diameter_endpoints; }
    const vector<int> &get_path() const { return diameter_path; }
};

// 树的重心 - 删除后使剩余连通分量最小的节点
// 使用场景: 树的分解、平衡划分、优化问题
// 时间复杂度: O(n)
// 典型应用: 树分治、负载均衡、网络设计
class TreeCentroid
{
private:
    const Tree &tree;
    vector<int> subtree_size;
    vector<bool> removed;
    int n;

    int dfs_size(int u, int parent)
    {
        subtree_size[u] = 1;
        for (int v : tree.adj[u]) {
            if (v != parent && !removed[v]) {
                subtree_size[u] += dfs_size(v, u);
            }
        }
        return subtree_size[u];
    }

    int find_centroid(int u, int parent, int tree_size)
    {
        for (int v : tree.adj[u]) {
            if (v != parent && !removed[v] && subtree_size[v] > tree_size / 2) {
                return find_centroid(v, u, tree_size);
            }
        }
        return u;
    }

public:
    TreeCentroid(const Tree &t) : tree(t), n(t.n), subtree_size(n), removed(n, false) {}

    int get_centroid(int root = 0)
    {
        int tree_size = dfs_size(root, -1);
        return find_centroid(root, -1, tree_size);
    }

    // 分治：递归处理以重心为根的子树
    template <typename ProcessFunc>
    void centroid_decomposition(int u, ProcessFunc process)
    {
        int centroid = get_centroid(u);
        removed[centroid] = true;

        // 处理当前重心
        process(centroid);

        // 递归处理子树
        for (int v : tree.adj[centroid]) {
            if (!removed[v]) {
                centroid_decomposition(v, process);
            }
        }
    }
};

// 树上路径查询 - 支持路径上的各种操作
// 使用场景: 路径统计、路径修改、树上数据结构
// 时间复杂度: 依赖具体操作，通常O(log n)到O(n)
// 典型应用: 树链剖分、重链剖分、树上莫队
class TreePath
{
private:
    const Tree &tree;
    LCA lca_solver;

public:
    TreePath(const Tree &t) : tree(t), lca_solver(t) {}

    // 获取u到v的路径上所有节点
    vector<int> get_path(int u, int v)
    {
        int lca_node = lca_solver.lca(u, v);
        vector<int> path1, path2;

        // u到lca的路径
        int curr = u;
        while (curr != lca_node) {
            path1.push_back(curr);
            curr = tree.parent[curr];
        }
        path1.push_back(lca_node);

        // v到lca的路径
        curr = v;
        while (curr != lca_node) {
            path2.push_back(curr);
            curr = tree.parent[curr];
        }

        // 合并路径
        vector<int> full_path = path1;
        for (int i = path2.size() - 1; i >= 0; i--) {
            full_path.push_back(path2[i]);
        }

        return full_path;
    }

    // 计算路径上的权值和
    ll path_sum(int u, int v)
    {
        int lca_node = lca_solver.lca(u, v);
        return tree.dist_from_root[u] + tree.dist_from_root[v] - 2 * tree.dist_from_root[lca_node];
    }

    // 检查点w是否在u到v的路径上
    bool on_path(int u, int v, int w)
    {
        int lca_uv = lca_solver.lca(u, v);
        int lca_uw = lca_solver.lca(u, w);
        int lca_vw = lca_solver.lca(v, w);

        return (lca_uw == w && lca_vw == lca_uv) || (lca_vw == w && lca_uw == lca_uv);
    }
};

// 换根DP - 通过换根实现全局最优解
// 使用场景: 树上优化问题、全局统计、动态规划
// 时间复杂度: O(n)
// 典型应用: 树的中心、最优根、距离和
template <typename T>
class RerootDP
{
private:
    const Tree &tree;
    vector<T> dp_down, dp_up; // down: 子树答案, up: 换根后答案
    int n;

public:
    RerootDP(const Tree &t) : tree(t), n(t.n), dp_down(n), dp_up(n) {}

    // 第一次DFS：计算每个节点作为子树根的答案
    template <typename DownFunc>
    void dfs_down(int u, int parent, DownFunc down_func)
    {
        vector<T> child_values;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                dfs_down(v, u, down_func);
                child_values.push_back(dp_down[v]);
            }
        }
        dp_down[u] = down_func(u, child_values);
    }

    // 第二次DFS：通过换根计算每个节点作为全树根的答案
    template <typename UpFunc>
    void dfs_up(int u, int parent, T parent_value, UpFunc up_func)
    {
        dp_up[u] = up_func(u, dp_down[u], parent_value);

        vector<T> child_values;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                child_values.push_back(dp_down[v]);
            }
        }

        int child_idx = 0;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                // 计算除了v子树外的其他部分对v的贡献
                vector<T> other_children;
                for (int i = 0; i < child_values.size(); i++) {
                    if (i != child_idx) {
                        other_children.push_back(child_values[i]);
                    }
                }

                T contribution = up_func(u, T{}, parent_value, other_children);
                dfs_up(v, u, contribution, up_func);
                child_idx++;
            }
        }
    }

    const vector<T> &get_down_results() const { return dp_down; }
    const vector<T> &get_up_results() const { return dp_up; }
};

// 树的DFS时间戳 - 将树转化为区间问题
// 使用场景: 子树查询、树上数据结构、区间操作
// 时间复杂度: O(n)
// 典型应用: 子树修改、祖先查询、树状数组
class TreeDFSOrder
{
private:
    const Tree &tree;
    vector<int> in_time, out_time, dfs_order;
    int timer;

public:
    TreeDFSOrder(const Tree &t) : tree(t), in_time(t.n), out_time(t.n), timer(0)
    {
        dfs(tree.root, -1);
    }

    void dfs(int u, int parent)
    {
        in_time[u] = timer++;
        dfs_order.push_back(u);

        for (int v : tree.adj[u]) {
            if (v != parent) {
                dfs(v, u);
            }
        }

        out_time[u] = timer - 1;
    }

    // 检查u是否是v的祖先
    bool is_ancestor(int u, int v)
    {
        return in_time[u] <= in_time[v] && out_time[v] <= out_time[u];
    }

    // 获取子树u对应的区间
    pair<int, int> subtree_range(int u)
    {
        return {in_time[u], out_time[u]};
    }

    const vector<int> &get_dfs_order() const { return dfs_order; }
    const vector<int> &get_in_time() const { return in_time; }
    const vector<int> &get_out_time() const { return out_time; }
};

// 树的拓扑排序 - 适用于有向无环图(DAG)
// 使用场景: 依赖关系处理、任务调度、编译顺序
// 时间复杂度: O(V + E)
// 典型应用: 课程安排、项目依赖、数据流分析
vector<int> tree_topological_sort(const vector<vector<int>> &adj, int n)
{
    vector<int> in_degree(n, 0);
    for (int u = 0; u < n; u++) {
        for (int v : adj[u]) {
            in_degree[v]++;
        }
    }

    queue<int> q;
    for (int i = 0; i < n; i++) {
        if (in_degree[i] == 0) {
            q.push(i);
        }
    }

    vector<int> result;
    while (!q.empty()) {
        int u = q.front();
        q.pop();
        result.push_back(u);

        for (int v : adj[u]) {
            in_degree[v]--;
            if (in_degree[v] == 0) {
                q.push(v);
            }
        }
    }

    return result.size() == n ? result : vector<int>(); // 返回空表示有环
}

// 树算法工具函数集合
namespace TreeUtils
{
// 计算子树大小
vector<int> calculate_subtree_sizes(const Tree &tree)
{
    vector<int> subtree_size(tree.n);

    function<int(int, int)> dfs = [&](int u, int parent) -> int {
        subtree_size[u] = 1;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                subtree_size[u] += dfs(v, u);
            }
        }
        return subtree_size[u];
    };

    dfs(tree.root, -1);
    return subtree_size;
}

// 计算每个节点到根的距离
vector<ll> calculate_distances_from_root(const Tree &tree)
{
    vector<ll> dist(tree.n, 0);

    function<void(int, int)> dfs = [&](int u, int parent) {
        for (auto [v, w] : tree.weighted_adj[u]) {
            if (v != parent) {
                dist[v] = dist[u] + w;
                dfs(v, u);
            }
        }
    };

    dfs(tree.root, -1);
    return dist;
}

// 计算树的高度
int calculate_tree_height(const Tree &tree)
{
    int max_height = 0;

    function<int(int, int)> dfs = [&](int u, int parent) -> int {
        int height = 0;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                height = max(height, 1 + dfs(v, u));
            }
        }
        max_height = max(max_height, height);
        return height;
    };

    dfs(tree.root, -1);
    return max_height;
}

// 找到树的所有叶子节点
vector<int> find_leaves(const Tree &tree)
{
    vector<int> leaves;
    for (int i = 0; i < tree.n; i++) {
        if (i != tree.root && tree.adj[i].size() == 1) {
            leaves.push_back(i);
        }
    }
    return leaves;
}

// 检查是否为二叉树
bool is_binary_tree(const Tree &tree)
{
    for (int i = 0; i < tree.n; i++) {
        int children = 0;
        for (int v : tree.adj[i]) {
            if (tree.parent[i] != v) {
                children++;
            }
        }
        if (children > 2) {
            return false;
        }
    }
    return true;
}

// 计算树的重心（可能有1-2个）
vector<int> find_centroids(const Tree &tree)
{
    vector<int> subtree_size(tree.n);
    vector<int> centroids;

    function<int(int, int)> dfs_size = [&](int u, int parent) -> int {
        subtree_size[u] = 1;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                subtree_size[u] += dfs_size(v, u);
            }
        }
        return subtree_size[u];
    };

    dfs_size(tree.root, -1);

    function<void(int, int)> dfs_centroid = [&](int u, int parent) {
        bool is_centroid = true;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                if (subtree_size[v] > tree.n / 2) {
                    is_centroid = false;
                }
                dfs_centroid(v, u);
            }
        }
        if (parent != -1 && tree.n - subtree_size[u] > tree.n / 2) {
            is_centroid = false;
        }
        if (is_centroid) {
            centroids.push_back(u);
        }
    };

    dfs_centroid(tree.root, -1);
    return centroids;
}

// 树的同构检测（简化版）
string tree_hash(const Tree &tree, int root)
{
    function<string(int, int)> dfs = [&](int u, int parent) -> string {
        vector<string> child_hashes;
        for (int v : tree.adj[u]) {
            if (v != parent) {
                child_hashes.push_back(dfs(v, u));
            }
        }
        sort(child_hashes.begin(), child_hashes.end());

        string result = "(";
        for (const string &child : child_hashes) {
            result += child;
        }
        result += ")";
        return result;
    };

    return dfs(root, -1);
}

// 计算所有节点对之间的距离
vector<vector<ll>> all_pairs_distances(const Tree &tree)
{
    vector<vector<ll>> dist(tree.n, vector<ll>(tree.n, 0));
    LCA lca_solver(tree);

    for (int i = 0; i < tree.n; i++) {
        for (int j = i + 1; j < tree.n; j++) {
            int lca_node = lca_solver.lca(i, j);
            ll distance = tree.dist_from_root[i] + tree.dist_from_root[j] - 2 * tree.dist_from_root[lca_node];
            dist[i][j] = dist[j][i] = distance;
        }
    }

    return dist;
}
} // namespace TreeUtils

// 树上启发式合并 (DSU on Tree) - 高效处理子树查询
// 使用场景: 子树内的复杂统计、颜色计数、路径查询
// 时间复杂度: O(n log n)
// 典型应用: 子树颜色数、子树内最值、复杂子树查询
template <typename T>
class DSUOnTree
{
private:
    const Tree &tree;
    vector<int> subtree_size, heavy_child;
    vector<T> answer;

    void calculate_heavy_children()
    {
        subtree_size.assign(tree.n, 0);
        heavy_child.assign(tree.n, -1);

        function<int(int, int)> dfs = [&](int u, int parent) -> int {
            subtree_size[u] = 1;
            int max_child_size = 0;

            for (int v : tree.adj[u]) {
                if (v != parent) {
                    int child_size = dfs(v, u);
                    subtree_size[u] += child_size;

                    if (child_size > max_child_size) {
                        max_child_size = child_size;
                        heavy_child[u] = v;
                    }
                }
            }

            return subtree_size[u];
        };

        dfs(tree.root, -1);
    }

public:
    DSUOnTree(const Tree &t) : tree(t), answer(t.n)
    {
        calculate_heavy_children();
    }

    // 启发式合并主函数
    template <typename AddFunc, typename RemoveFunc, typename QueryFunc>
    void solve(int u, int parent, bool keep, AddFunc add, RemoveFunc remove, QueryFunc query)
    {
        // 先处理轻儿子
        for (int v : tree.adj[u]) {
            if (v != parent && v != heavy_child[u]) {
                solve(v, u, false, add, remove, query);
            }
        }

        // 处理重儿子
        if (heavy_child[u] != -1) {
            solve(heavy_child[u], u, true, add, remove, query);
        }

        // 添加轻儿子的贡献
        for (int v : tree.adj[u]) {
            if (v != parent && v != heavy_child[u]) {
                function<void(int, int)> add_subtree = [&](int node, int p) {
                    add(node);
                    for (int child : tree.adj[node]) {
                        if (child != p) {
                            add_subtree(child, node);
                        }
                    }
                };
                add_subtree(v, u);
            }
        }

        // 添加当前节点
        add(u);

        // 计算答案
        answer[u] = query(u);

        // 如果不保留，清除所有贡献
        if (!keep) {
            function<void(int, int)> remove_subtree = [&](int node, int p) {
                remove(node);
                for (int child : tree.adj[node]) {
                    if (child != p) {
                        remove_subtree(child, node);
                    }
                }
            };
            remove_subtree(u, parent);
        }
    }

    const vector<T> &get_answers() const { return answer; }
};
